#pragma once
#include <iostream>
#include <cassert>
template<class k, class value>
struct AVLNode
{
	std::pair<k, value> _kv;
	AVLNode<k, value>* _left;
	AVLNode<k, value>* _right;
	AVLNode<k, value>* _parent;
	int _bf;

	AVLNode(const std::pair<k, value>& kv)
		:_kv(kv), _left(nullptr), _right(nullptr), _parent(nullptr), _bf(0)
	{}
};

template<class k, class value>
class AVLTree
{
	typedef AVLNode<k, value> Node;
	typedef std::pair<k, value> PKV;
private:
	Node* _root = nullptr;
private:
	void _InOrder(Node* root);
	int Height(Node* root);
	bool _IsBalance(Node* root);
	void AVL_Rotate_LL(Node* parent);
	void AVL_Rotate_RR(Node* parent);
	void AVL_Rotate_LR(Node* parent);
	void AVL_Rotate_RL(Node* parent);
public:
	bool IsBalance();
	k root();
	void InOrder();
	bool Insert(const PKV& kv);
};

template<class k, class value>
k AVLTree<k, value>::root()
{
	if (root == nullptr) return k();
	return _root->_kv.first;
}
template<class k, class value>
int AVLTree<k, value>::Height(Node* root)
{
	if (root == nullptr) return 0;
	int left = Height(root->_left);
	int right = Height(root->_right);
	return left > right ? left + 1 : right + 1;
}

template<class k, class value>
bool AVLTree<k, value>::_IsBalance(Node* root)
{
	if (root == nullptr) return true;
	int left = Height(root->_left);
	int right = Height(root->_right);
	return abs(left - right) < 2 && _IsBalance(root->_left) && _IsBalance(root->_right);
}

template<class k, class value>
bool AVLTree<k, value>::IsBalance()
{
	if (_root == nullptr) return true;
	return _IsBalance(_root);
}

template<class k, class value>
void AVLTree<k, value>::_InOrder(Node* root)
{
	if (root == nullptr) return;
	_InOrder(root->_left);
	std::cout << root->_kv.first << ' ';
	_InOrder(root->_right);
}

template<class k, class value>
void AVLTree<k, value>::InOrder()
{
	_InOrder(_root);
	std::cout << std::endl;
}

template<class k, class value>
void AVLTree<k, value>::AVL_Rotate_LL(Node* parent)
{
	Node* cur = parent->_left;
	//1. 孩子的右给父亲的左
	if (cur->_right)
	{
		cur->_right->_parent = parent;
	}
	parent->_left = cur->_right;
	//2. 孩子的右指向父亲
	cur->_right = parent;
	//3. 更改cur和parent的父亲 -- 注意pparent
	Node* pparent = parent->_parent;
	cur->_parent = pparent;
	parent->_parent = cur;
	if (pparent)
	{
		if (pparent->_left == parent) pparent->_left = cur;
		else pparent->_right = cur;
	}
	else
	{
		_root = cur;
	}
	//4. 调整平衡因子
	cur->_bf = parent->_bf = 0;
}
template<class k, class value>
void AVLTree<k, value>::AVL_Rotate_RR(Node* parent)
{
	Node* cur = parent->_right;
	//1. 孩子的左给父亲的右
	if (cur->_left)
	{
		cur->_left->_parent = parent;
	}
	parent->_right = cur->_left;
	//2. 孩子的左指向父亲
	cur->_left = parent;
	//3. 更改cur和parent的父亲 -- 注意pparent
	Node* pparent = parent->_parent;
	cur->_parent = pparent;
	parent->_parent = cur;
	if (pparent)
	{
		if (pparent->_left == parent) pparent->_left = cur;
		else pparent->_right = cur;
	}
	else 
	{
		_root = cur;
	}
	//4. 调整平衡因子
	cur->_bf = parent->_bf = 0;
}
template<class k, class value>
void AVLTree<k, value>::AVL_Rotate_LR(Node* parent)
{
	//A = parent, B = cur, C = cur->_right
	Node* A = parent;
	Node* B = parent->_left;
	Node* C = B->_right;
	int bf = C->_bf;
	AVL_Rotate_RR(B);
	AVL_Rotate_LL(A);
	if (bf == 0)
	{
		A->_bf = B->_bf = C->_bf = 0;
	}
	else if (bf == 1)
	{
		B->_bf = -1;
		C->_bf = A->_bf = 0;
	}
	else if (bf == -1)
	{
		B->_bf = C->_bf = 0;
		A->_bf = 1;
	}
}
template<class k, class value>
void AVLTree<k, value>::AVL_Rotate_RL(Node* parent)
{
	//A = parent, B = cur, C = cur->_left
	Node* A = parent;
	Node* B = parent->_right;
	Node* C = B->_left;
	int bf = C->_bf;
	AVL_Rotate_LL(B);
	AVL_Rotate_RR(A);
	//经过2次旋转，C的左给B的右，C的右给A的左
	//C的bf有3种情况
	if (bf == 0)
	{
		A->_bf = B->_bf = C->_bf = 0;
	}
	else if (bf == 1)
	{
		A->_bf = -1;
		B->_bf = C->_bf = 0;
	}
	else if (bf == -1)
	{
		A->_bf = C->_bf = 0;
		B->_bf = 1;
	}
}
template<class k, class value>
bool AVLTree<k, value>::Insert(const PKV& kv)
{
	if (_root == nullptr)
	{
		_root = new Node(kv);
		return true;
	}
	//1. 找插入位置
	Node* cur = _root;
	Node* parent = cur->_parent;
	while (cur)
	{
		if (kv.first > cur->_kv.first) parent = cur, cur = cur->_right;
		else if (kv.first < cur->_kv.first) parent = cur, cur = cur->_left;
		else return false;
	}
	cur = new Node(kv);
	if (kv.first > parent->_kv.first)
	{
		parent->_right = cur;
	}
	else if (kv.first < parent->_kv.first)
	{
		parent->_left = cur;
	}
	else
	{
		assert(false);
	}
	cur->_parent = parent;
	//2. 调整平衡因子
	while (parent)
	{
		//2.1 调整平衡因子
		if (parent->_left == cur)
		{
			parent->_bf++;
		}
		else
		{
			parent->_bf--;
		}
		//2.2 判断平衡因子
		if (parent->_bf == 0)
		{
			//爷爷一定平衡
			break;
		}
		else if (parent->_bf == 1 || parent->_bf == -1)
		{
			//爷爷可能失衡, 继续调整
			cur = parent;
			parent = cur->_parent;
		}
		else if (parent->_bf == 2 || parent->_bf == -2)
		{
			if (parent->_bf == 2 && cur->_bf == 1)
			{
				//LL型 -- 调整后平衡因子 = 0， 不用继续向上调整
				AVL_Rotate_LL(parent);
			}
			else if (parent->_bf == -2 && cur->_bf == -1)
			{
				//RR型
				AVL_Rotate_RR(parent);
			}
			else if (parent->_bf == 2 && cur->_bf == -1)
			{
				AVL_Rotate_LR(parent);
			}
			else if (parent->_bf == -2 && cur->_bf == 1)
			{
				AVL_Rotate_RL(parent);
			}
			break;
		}
		else
		{
			assert(false);
		}
	}
	return true;
}